import gradio as gr
import numpy as np
from PIL import Image
from ocr_process.cls.clas_pipeline import ClasPipeline

clas_predict = ClasPipeline(model_path=r"D:\ocr\model\multi_class\multi_class_1.6\inference.onnx")


def predict(image: Image.Image):
    image_data = np.array(image)
    input_data = {"original_image": image_data, "image": image_data.copy()}
    result = clas_predict(input_data)
    print(result)
    return result


def gradio_interface(image):
    return predict(image)


iface = gr.Interface(
    fn=gradio_interface,
    inputs="image",
    outputs="text",
    title="图片分类系统",
    description="上传一张图片，模型将预测其类别。支持荣钢收发货,天津港、荣钢发货(包含多张榜单)）"
)

if __name__ == '__main__':
    iface.launch(server_name="0.0.0.0", share=True)
